// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/expression/complex_view_types.h>
#include <pollux/functions/udf.h>

namespace kumo::pollux::functions {
    template<typename TExec>
    struct MapTopNFunction {
        POLLUX_DEFINE_FUNCTION_TYPES(TExec);

        template<typename It>
        struct Compare {
            bool operator()(const It &l, const It &r) const {
                static const CompareFlags flags{
                    false /*nullsFirst*/,
                    true /*ascending*/,
                    false /*equalsOnly*/,
                    CompareFlags::NullHandlingMode::kNullAsIndeterminate
                };

                if (l->second.has_value() && r->second.has_value()) {
                    auto comp = l->second.value().compare(r->second.value(), flags);

                    if (MELON_UNLIKELY(comp == 0)) {
                        return l->first.compare(r->first, flags) > 0;
                    }

                    return comp > 0;
                } else if (MELON_UNLIKELY(
                    !l->second.has_value() && !r->second.has_value())) {
                    return l->first.compare(r->first, flags) > 0;
                }

                return l->second.has_value();
            }
        };

        void call(
            out_type<Map<Orderable<T1>, Orderable<T2> > > &out,
            const arg_type<Map<Orderable<T1>, Orderable<T2> > > &inputMap,
            int64_t n) {
            POLLUX_USER_CHECK_GE(n, 0, "n must be greater than or equal to 0");

            if (n == 0) {
                return;
            }

            using It = typename arg_type<Map<Orderable<T1>, Orderable<T2> > >::Iterator;
            std::vector<It> vec;
            vec.reserve(inputMap.size());
            for (auto it = inputMap.begin(); it != inputMap.end(); ++it) {
                vec.push_back(it);
            }

            Compare<It> comparator;
            if (n >= inputMap.size()) {
                // Use std::sort to handle nulls when n >= inputMap.size(), as
                // std::nth_element doesn't invoke the comparator in this case.
                std::sort(vec.begin(), vec.end(), comparator);
                out.copy_from(inputMap);
                return;
            }

            std::nth_element(vec.begin(), vec.begin() + n, vec.end(), comparator);
            for (auto it = vec.begin(); it != vec.begin() + n; ++it) {
                const auto &ele = **it;
                if (!ele.second.has_value()) {
                    auto &keyWriter = out.add_null();
                    keyWriter.copy_from(ele.first);
                } else {
                    auto [keyWriter, valueWriter] = out.add_item();
                    keyWriter.copy_from(ele.first);
                    valueWriter.copy_from(ele.second.value());
                }
            }
        }
    };
} // namespace kumo::pollux::functions
